package server

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"time"

	"yx.io/qwq-ai/ai"
	"yx.io/qwq-ai/myComponent"
)

func sseHandler(w http.ResponseWriter, r *http.Request) {
	// 设置 SSE 所需的 HTTP 头
	w.Header().Set("Content-Type", "text/event-stream;charset=utf-8")
	w.Header().Set("Cache-Control", "no-cache")
	w.Header().Set("Connection", "keep-alive")
	w.Header().Set("Access-Control-Allow-Origin", "*")

	flusher, ok := w.(http.Flusher)
	if !ok {
		log.Println("不支持 SSE")
		return
	}

	// 获取用户的问题
	q := r.URL.Query().Get("q")
	if q == "" {
		resultJson := SSEMessage{
			Data:     "请输入问题",
			DateTime: time.Now().Format("2006-01-02 15:04:05"),
		}
		resultStr, _ := json.Marshal(resultJson)
		fmt.Fprintf(w, "data: %s\n\n", resultStr)
		// 发送结束事件
		fmt.Fprintf(w, "%s", "event: end\ndata: \n\n")
		flusher.Flush()
		return
	}

	sidCookie, err := r.Cookie(GSESSION_ID)
	var sessionId string
	var messageList myComponent.MessageList
	if err == http.ErrNoCookie || sidCookie == nil { // 用户浏览器中没有sessionId的cookie
		sessionId = myComponent.GenerateSessionId()
		messageList.SessionId = sessionId
	} else {
		sessionId = sidCookie.Value
	}
	if messageList, ok = myComponent.SessionContainer[sessionId]; !ok { // 服务器上没有这个sessionId的消息列表
		sessionId = myComponent.GenerateSessionId()
		messageList.SessionId = sessionId
	}

	if q == "clear" {
		// 清除历史
		myComponent.SessionContainer[sessionId] = myComponent.MessageList{
			SessionId: sessionId,
			Messages:  make([]myComponent.Message, 0),
		}
		resultJson := SSEMessage{
			Data:     "历史已清除",
			DateTime: time.Now().Format("2006-01-02 15:04:05"),
		}
		resultStr, _ := json.Marshal(resultJson)
		fmt.Fprintf(w, "data: %s\n\n", resultStr)
		// 发送结束事件
		fmt.Fprintf(w, "%s", "event: end\ndata: \n\n")
		flusher.Flush()
		return
	}

	// 重新写回cookie
	cookie := &http.Cookie{
		Name:   GSESSION_ID,
		Value:  sessionId,
		MaxAge: 60 * 60 * 24,
		Path:   "/",
	}
	http.SetCookie(w, cookie)

	ch := make(chan string)
	// 启动 goroutine 处理问题
	go func() {
		// 定义管道，
		ai.QwqStreamBySessionId(q, ch, sessionId)
	}()

	answer := ""
	for res := range ch {
		answer += res
		resultJson := SSEMessage{
			Data:     res,
			DateTime: time.Now().Format("2006-01-02 15:04:05"),
		}
		resultStr, _ := json.Marshal(resultJson)
		fmt.Fprintf(w, "data: %s\n\n", resultStr)
		flusher.Flush()
	}
	// 发送结束事件
	fmt.Fprintf(w, "%s", "event: end\ndata: \n\n")
	flusher.Flush()

	ip := r.Context().Value(IPKey)
	log.Printf("ip:%s question:%s answer: %s", ip, q, answer)

	// 写回sessionContainer
	messageList.Messages = append(messageList.Messages, myComponent.Message{
		Role:    "user",
		Content: q,
	}, myComponent.Message{
		Role:    "system",
		Content: answer,
	})

	messageList.ExpireTime = time.Now().Add(24 * time.Hour) // 更新过期时间为24小时后
	myComponent.SessionContainer[messageList.SessionId] = messageList

}

// 给前端封装的消息结构体
type SSEMessage struct {
	Data     string `json:"data"`
	DateTime string `json:"dateTime"`
}
